import datetime

import mplfinance as mpf
import numpy as np
import pandas as pd
import talib as tb
import tushare as ts


def get_stock(_stock, _pro):

    today = datetime.datetime.today()
    startday = today+datetime.timedelta(days=-365)
    today = today.strftime('%Y%m%d')
    startday = startday.strftime('%Y%m%d')
    # 获取股票数据
    onestock_df = _pro.daily(
        ts_code=_stock, start_date=startday, end_date=today)
    # 将trade_date转换为时间格式
    onestock_df['trade_date'] = pd.to_datetime(onestock_df['trade_date'])
    # 倒序排列 iloc[::-1]
    onestock_df = onestock_df.iloc[::-1]
    # 将列vol改为volume
    onestock_df = onestock_df.rename(columns={'vol': 'volume'})
    # 保存为csv文件，不保存索引
    onestock_df.to_csv('%s.csv' % _stock, index=False)
    # 读取csv文件,将trade_date作为行索引
    onestock_df = pd.read_csv('%s.csv' % _stock, index_col=1)

    # 将索引转为时间格式
    onestock_df.index = pd.to_datetime(onestock_df.index)
    return onestock_df


def trade_signal(_onestock_df, macd):
    """
    计算交易信号，当股价低于布林下轨时买入，高于上轨时卖出
    :param code: 股票代码
    :return: 包含交易信号的DataFrame
    """
    # 新建一个DataFrame,以data的index为index
    df = pd.DataFrame(index=_onestock_df.index)

    df[['price', 'upper', 'lower']] = _onestock_df[['close', 'high', 'low']]

    # 初始化订单状态为0
    df['orders'] = 0

    # 初始仓位为0
    position = 0

    # 遍历数据表
    for i in range(len(_onestock_df)):
        # 当买入信号为True且仓位为0时买入
        if macd.Bar.values[i-1] < 0 and macd.Bar.values[i] > 0 and position == 0:
            # 买入指令为1
            df.orders.values[i] = 1

            position = 1

        # 当卖出信号为True且仓位为1时卖出
        if macd.Bar.values[i-1] > 0 and macd.Bar.values[i] < 0 and position == 1:
            # 卖出指令为-1
            df.orders.values[i] = -1

            position = 0

    return df


# 计算收益
def calculate_income(_money, _onestock_df, macd, _print):
    """
    计算收益
    :param code: 股票代码
    :return: 计算后的DataFrame
    """

    data = trade_signal(_onestock_df, macd)
    # 初始最大可持仓量,可用资金为100000元
    max_position = 0
    money = _money * 10000  # 100万
    shiZhi = 0
    ciShu = 0

    # 初始持仓股票手数、可用资金
    num = 0
    first_buy_price = 0
    end_sell_price = 0
    first_buy_data = 0
    end_sell_data = 0
    price_rate = 0

    for i in range(len(data)):
        # 当买卖指令为1
        if data.orders.values[i] == 1:
            # 以下轨价全仓买入
            max_position = money // (data.price.values[i] * 100)
            # 可用资金减去买入股票所花费的资金
            money = money - max_position * data.price.values[i] * 100
            # 记录第一次买入时价格 和 日期
            if first_buy_price == 0:
                first_buy_price = data.price.values[i]
                first_buy_data = data.index
            if _print == 1:
                print("买入日期", data.index[i], "买入价格", data.price.values[i], "手数",
                      max_position, "市值",  max_position * data.price.values[i] * 100)
        # 当买卖指令为-1
        elif data.orders.values[i] == -1:
            # 可用资金加上卖出股票所得到的资金

            money = money + max_position * data.price.values[i] * 100
            shiZhi = money
            ciShu = num
            # 持仓清零
            num = num + 1

            # 记录最后一次卖出时价格
            if end_sell_price == 0:
                end_sell_price = data.price.values[i]
                end_sell_data = data.index
            if _print == 1:
                print("卖出日期", data.index[i], "卖出价格", data.price.values[i],
                      "手数", max_position, "余额",  money, "卖出次数 ", num)

            max_position = 0
        # 第一次买入和最后一次卖出价格变化率
    price_rate = round(
        (end_sell_price - first_buy_price)/first_buy_price*100, 3)

    return shiZhi, ciShu, price_rate, end_sell_price


def simTradePlot(_onestock_df, macd, redBar, greenBar):

    # 设置mplfinance的蜡烛颜色
    # up为阳线颜色
    # down为阴线颜色
    my_color = mpf.make_marketcolors(
        up='red',
        down='limegreen',
        edge='inherit',
        wick='inherit',
        volume='inherit')

    # 设置图形风格
    # figcolor:设置图表的背景色
    # y_on_right:设置y轴位置是否在右
    # gridaxis:设置网格线位置
    # gridstyle:设置网格线线型
    # gridcolor:设置网格线颜色
    my_style = mpf.make_mpf_style(
        marketcolors=my_color,
        figcolor='#EEEEEE',
        y_on_right=False,
        gridaxis='both',
        gridstyle='-.',
        gridcolor='#E1E1E1')

    # 设置基本参数
    # type:绘制图形的类型，有candle, renko, ohlc, line等
    # 此处选择candle,即K线图
    # mav(moving average):均线类型,此处设置5,10,30日线
    # volume:布尔类型，设置是否显示成交量，默认False
    # title:设置标题
    # y_label:设置纵轴主标题
    # y_label_lower:设置成交量图一栏的标题
    # figratio:设置图形纵横比
    # figscale:设置图形尺寸(数值越大图像质量越高)
    # datetime_format:设置日期显示格式
    # xrotation:设置x坐标的转角度
    kwargs = dict(
        type='candle',
        mav=(5, 10, 30),
        volume=True,
        title='%s' % (_onestock_df.iloc[0, 0]),
        ylabel='Price',
        ylabel_lower='Volume',
        figratio=(1200/72, 480/60),
        figscale=3,
        datetime_format='%Y-%m-%d',
        xrotation=15)

    # RSI
    rsi = tb.RSI(_onestock_df.open, timeperiod=12)
    # rsi=pd.DataFrame(rsi_df,columns=['0'])

    # 设置配图
    add_plot = [
        mpf.make_addplot(redBar, type='bar', panel=2,
                         ylabel='MACD', color='red'),
        mpf.make_addplot(greenBar, type='bar',
                         panel=2, color='limegreen'),
        mpf.make_addplot(macd['date'], panel=2, color='orangered'),
        mpf.make_addplot(macd['signal'], panel=2, color='limegreen'),

    ]

    mpf.plot(_onestock_df, **kwargs, addplot=add_plot, style=my_style)


def cal_macd(_onestock_df):
    # MACD

    macd = pd.DataFrame(index=_onestock_df.index)

    macd_date, macdsignal, macd_bar = tb.MACD(
        _onestock_df['close'], fastperiod=12, slowperiod=26, signalperiod=9)

    macd['date'] = pd.DataFrame(macd_date, columns=['0'])
    macd['signal'] = pd.DataFrame(macdsignal, columns=['0'])
    macd['Bar'] = pd.DataFrame(macd_bar, columns=['0'])

    bar_red = np.where(macd_bar > 0, 2 * macd_bar, 0)  # 绘制BAR>0 柱状图
    bar_green = np.where(macd_bar < 0, 2 * macd_bar, 0)  # 绘制BAR<0 柱状图

    redBar = pd.DataFrame(bar_red, columns=['0'])
    greenBar = pd.DataFrame(bar_green, columns=['0'])
    return macd, redBar, greenBar


def sim_trade(_money, stock_code, pro, _printf, _plot):
    onestock_df = get_stock(stock_code, pro)

    macd, redBar, greenBar = cal_macd(onestock_df)

    data, num, price_rate, end_sell_price = calculate_income(
        _money, onestock_df, macd, _printf)
    if _plot == 1:
        simTradePlot(onestock_df, macd, redBar, greenBar)
    return data, num, price_rate, end_sell_price


# token = '722a97e6b7907534837b7dab985a012166179358cc659fe061231317'
# ts.set_token(token)
# pro = ts.pro_api()
# sim_trade('000543.SZ', pro)
